package com.csust.automaticteach.app.rag;


import com.csust.automaticteach.constant.ChromaConstant;
import com.fasterxml.jackson.databind.ObjectMapper;
import lombok.Data;
import org.springframework.ai.chroma.vectorstore.ChromaApi;
import org.springframework.ai.chroma.vectorstore.ChromaVectorStore;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.boot.context.properties.ConfigurationProperties;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.http.client.SimpleClientHttpRequestFactory;
import org.springframework.web.client.RestClient;

import static com.csust.automaticteach.constant.ChromaConstant.*;

@Configuration
@ConfigurationProperties(prefix = "vectorstore.chroma")
@Data
public class ChromaVectorStoreConfig {

    private String baseUrl;
    @Bean
    public VectorStore chromaVectorStore(
            @Qualifier("dashscopeEmbeddingModel") EmbeddingModel embeddingModel,ChromaApi chromaApi) {
        return ChromaVectorStore.builder(chromaApi, embeddingModel)
                .tenantName(CHROMA_TENANT) // default: SpringAiTenant
                .databaseName(CHROMA_DATABASE) // default: SpringAiDatabase
                .collectionName(CHROMA_COLLECTION)// default: SpringAiCollection
                .initializeSchema(true)
                .build();
    }
    @Bean
    public RestClient.Builder builder() {
        return RestClient.builder().requestFactory(new SimpleClientHttpRequestFactory());
    }


    @Bean
    public ChromaApi chromaApi(RestClient.Builder restClientBuilder) {
        ObjectMapper objectMapper = new ObjectMapper();
        return ChromaApi.builder()
                .baseUrl(baseUrl)
                .objectMapper(objectMapper)
                .restClientBuilder(restClientBuilder)
                .build();
    }
}
